import matplotlib.pyplot as plt
import matplotlib
import numpy as np
#
matplotlib.rc('xtick', labelsize=17)
matplotlib.rc('ytick', labelsize=17)
def smooth(data, weight=0.5):
    scalar = data
    last = scalar[0]
    smoothed = []
    for point in scalar:
        smoothed_val = last * weight + (1 - weight) * point
        smoothed.append(smoothed_val)
        last = smoothed_val

    return np.array(smoothed)
idx = 0
f = plt.figure(figsize=[10, 4])
# exp_list = [1, 3, 5]
exp_list = [["e7","e8","ee8","e9"],["e10","e11","ee11","e12"]]
for x in exp_list:
    ax = plt.subplot(1, 2, idx + 1)
    idx += 1
    acc = np.load('../save/result_save/{}/acc.npy'.format(x[0]), allow_pickle=True)  # avg
    acc1 = np.load('../save/result_save/{}/acc.npy'.format(x[1]), allow_pickle=True)  # fml
    acc2 = np.load('../save/result_save/{}/acc.npy'.format(x[2]), allow_pickle=True)  # fml
    acc3 = np.load('../save/result_save/{}/acc.npy'.format(x[3]), allow_pickle=True)  # fml
    # acc_iid = smooth(acc_iid)
    # acc_noniid = smooth(acc_noniid)
    plt.plot(range(len(acc)), acc, '--', linewidth=3.0, label='FedAvg',
             color="#EE442A")
    plt.plot(range(len(acc1)), acc1, linewidth=3.0, label=r"Fedprox,$\mu$=0.1",
             color="#5D35CF")
    plt.plot(range(len(acc2)), acc2, '--',linewidth=3.0, label=r"Fedprox,$\mu$=0.01",
             color="green")
    plt.plot(range(len(acc3)), acc3, linewidth=3.0, label='FML',
             color="#319FB3")
    plt.xlabel("# Rounds", fontsize=22)
    plt.xticks(fontsize=17)
    plt.yticks(fontsize=17)

    if idx == 1:
        plt.ylabel('Testing Accuracy', fontsize=18)
    if idx == 1:
        plt.title("MLP", fontsize=18)
    elif idx == 2:
        plt.title("LeNet5", fontsize=18)
    ax.tick_params(color='#dddddd')
    ax.spines['bottom'].set_color('#dddddd')
    ax.spines['top'].set_color('#dddddd')
    ax.spines['right'].set_color('#dddddd')
    ax.spines['left'].set_color('#dddddd')

    if idx == 2:
        plt.legend(fontsize=15)
    plt.tight_layout()

f.savefig("accuracy_reduction.pdf")
print("finish")